from .network_v2 import TactileDeltaEstimator, LayerUnfreezer, TactileDeltaTrainer
from .delta_inferencer import TactileDeltaMeter
import torch
import numpy as np
import os


class DeltaEstimator:
    """
    DeltaEstimator 类用于估计滑移距离 δ ∈ [0, 1]。
    它提供了基于触觉图像的滑移距离估计功能，并整合了网络的在线训练功能
    """

    def __init__(self, res_dir=None, online_training = False):
        """
        初始化 DeltaEstimator。
        """
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.model = TactileDeltaEstimator().to(self.device)
        self.online_training = online_training
        if self.online_training:
            self.model.train()
        else:
            self.model.eval()
        self.trainer = TactileDeltaTrainer(
            model=self.model,
            device=self.device,
            lr=1e-3,
            weight_decay=1e-5,
            loss_fn=torch.nn.SmoothL1Loss(beta=0.3),
            scheduler_type="plateau",
            scheduler_patience=5,
            scheduler_factor=0.5,
        )
        self.layer_unfreezer = LayerUnfreezer(self.model)
        self.inferencer = TactileDeltaMeter(model=self.model, device=self.device)

        current_dir = os.path.dirname(os.path.abspath(__file__))
        self.res_lib = os.path.join(current_dir, "resources")
        if res_dir is None:
            self.res_dir = self._get_newest_res_dir(self.res_lib)
        else:
            self.res_dir = os.path.join(self.res_dir, res_dir)
            if not self._is_valid_res_dir(self.res_dir):
                raise ValueError(f"Invalid resource path: {self.res_dir}")

        if self.res_dir is not None:
            print(f"[DeltaEstimator] Using resources from {self.res_dir}")
            self._load_resources(self.res_dir)
        else:
            print(f"[DeltaEstimator] Using random start network, no resources loaded.")

        self.working_dir = self._get_working_res_dir(self.res_dir)
        self.trainer.set_working_dir(self.working_dir)

    # Short cut for inferencer
    def data_to_delta(self, P, D, F):
        return self.inferencer.data_to_delta(P, D, F)

    def _get_newest_res_dir(self, res_lib):
        """
        获取指定目录下最新的资源文件路径。
        :param res_dir: 资源目录
        :return: 最新的资源文件夹路径
        """
        if not os.path.exists(res_lib):
            raise FileNotFoundError(f"Resource library does not exist: {res_lib}")

        res_idx = 0
        while True:
            res_dir = os.path.join(res_lib, f"res_{res_idx}")
            if not self._is_valid_res_dir(res_dir):
                res_idx -= 1  # 如果没有有效的资源目录，则返回上一个有效目录
                break
            res_idx += 1
        res_dir = os.path.join(res_lib, f"res_{res_idx}")
        if not self._is_valid_res_dir(res_dir):
            return None
        return res_dir

    def _get_working_res_dir(self, res_dir):
        """
        获取工作目录。
        :param res_dir: 资源目录
        :return: 工作目录，是资源目录的下一个目录
        """
        if res_dir is None:
            return os.path.join(self.res_lib, "res_0")

        res_idx = int(res_dir.split("_")[-1])
        working_dir = os.path.join(self.res_lib, f"res_{res_idx + 1}")
        if not os.path.exists(working_dir):
            os.makedirs(working_dir)

        return working_dir

    def _is_valid_res_dir(self, res_dir):
        """
        检查资源目录是否有效。
        :param res_dir: 资源目录
        :return: 是否有效
        """
        if not os.path.exists(res_dir):
            return False
        if not os.path.isdir(res_dir):
            return False
        if not os.path.exists(os.path.join(res_dir, "best_model.pth")):
            return False
        if not os.path.exists(os.path.join(res_dir, "mean.npy")):
            return False
        if not os.path.exists(os.path.join(res_dir, "std.npy")):
            return False
        if not os.path.exists(os.path.join(res_dir, "unfreezed_layers.npy")):
            return False
        if not os.path.exists(os.path.join(res_dir, "stable")):
            print("Warning: Resource directory has all the resource files but does not have the stable file.")
            return False
        return True

    def _load_resources(self, res_dir):
        model_data = torch.load(os.path.join(res_dir, "best_model.pth"), map_location=self.device, weights_only=False)
        self.model.load_state_dict(model_data["model_state_dict"])
        # self.trainer.optimizer.load_state_dict(model_data["optimizer_state_dict"])
        self.inferencer.mean = np.load(os.path.join(res_dir, "mean.npy"))
        self.inferencer.std = np.load(os.path.join(res_dir, "std.npy"))
        unfreezed_layers = np.load(os.path.join(res_dir, "unfreezed_layers.npy"), allow_pickle=True).item()
        self.layer_unfreezer.unfreeze_n_layers(unfreezed_layers)

    def save_mean_std(self, mean=None, std=None):
        """
        保存均值和标准差到资源目录。
        """
        if mean is not None:
            self.inferencer.mean = mean
        if std is not None:
            self.inferencer.std = std
        np.save(os.path.join(self.working_dir, "mean.npy"), self.inferencer.mean)
        np.save(os.path.join(self.working_dir, "std.npy"), self.inferencer.std)

    def save_unfreeze_layer(self):
        np.save(os.path.join(self.working_dir, "unfreezed_layers.npy"), self.layer_unfreezer._current_index)

if __name__ == "__main__":
    # test inference
    delta_estimator = DeltaEstimator()
    tacdata = torch.tensor(torch.rand((400,9))).to(delta_estimator.device)
    P,D,F = tacdata[:, :3], tacdata[:, 3:6], tacdata[:, 6:]
    # calculate the time used for inference
    st_time = torch.cuda.Event(enable_timing=True)
    ed_time = torch.cuda.Event(enable_timing=True)
    st_time.record()
    delta = delta_estimator.data_to_delta(P, D, F)
    ed_time.record()
    torch.cuda.synchronize()
    print(f"Delta: {delta}, Time: {st_time.elapsed_time(ed_time)} ms")